#!/usr/bin/env bash
set -xeuo pipefail

# ======================== 用户可修改区 =========================

project_name='DAPO'
exp_name='DAPO-Qwen2.5-3B-med'

# 数据与模型路径（请按需修改）
TRAIN_FILE='/home/zhangpinglu/data0/gy/code/fundus-reasoner-rl/experiments/dataset/RL_parquet/med_pub_rl_train_rmNormal1.parquet'
VAL_FILE='/home/zhangpinglu/data0/gy/code/fundus-reasoner-rl/experiments/dataset/RL_parquet/med_pub_rl_test.parquet'
MODEL_PATH='/home/zhangpinglu/data0/gy/code/fundus-reasoner-adaptive/experiments/ckpt/sft_stage3_cot'

train_prompt_bsz=16
gen_prompt_bsz=64
n_resp_per_prompt=16
train_prompt_mini_bsz=16

max_prompt_length=1024
max_response_length=4096

clip_ratio_low=0.2
clip_ratio_high=0.28

gpu_num=8
sp_size=8    # 与卡数一致
gen_tp=4
offload=True
use_dynamic_bsz=True

# overlong buffer（按需启用/禁用）
enable_overlong_buffer=True
overlong_buffer_len=1024
overlong_penalty_factor=1.0

# ======================== 不建议改动区 =========================

RUNTIME_ENV="./verl/trainer/runtime_env.yaml"  # 你的verl环境依赖包
CKPTS_DIR="./experiments/ckpts/${project_name}/${exp_name}"

# DeepSpeed/FSDP环境变量
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8   # 按你的服务器显卡编号分配

# 训练主程序
python3 -m recipe.med_rl.main_dapo \
    data.train_files="${TRAIN_FILE}" \
    data.val_files="${VAL_FILE}" \
    data.prompt_key=prompt \
    data.truncation='left' \
    data.max_prompt_length=${max_prompt_length} \
    data.max_response_length=${max_response_length} \
    data.gen_batch_size=${gen_prompt_bsz} \
    data.train_batch_size=${train_prompt_bsz} \
    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
    algorithm.adv_estimator=grpo \
    algorithm.use_kl_in_reward=False \
    algorithm.kl_ctrl.kl_coef=0.0 \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.kl_loss_coef=0.0 \
    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \
    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \
    actor_rollout_ref.actor.clip_ratio_c=10.0 \
    algorithm.filter_groups.enable=False \
    algorithm.filter_groups.max_num_gen_batches=20 \
    algorithm.filter_groups.metric=acc \
    actor_rollout_ref.rollout.name=vllm \
    actor_rollout_ref.model.use_remove_padding=True \
    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \
    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \
    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$((max_prompt_length + max_response_length)) \
    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$((max_prompt_length + max_response_length)) \
    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$((max_prompt_length + max_response_length)) \
    actor_rollout_ref.model.path="${MODEL_PATH}" \
    actor_rollout_ref.model.enable_gradient_checkpointing=True \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.optim.lr_warmup_steps=20 \
    actor_rollout_ref.actor.optim.weight_decay=0.1 \
    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \
    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \
    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \
    actor_rollout_ref.actor.entropy_coeff=0 \
    actor_rollout_ref.actor.grad_clip=1.0 \
    actor_rollout_ref.actor.loss_agg_mode=token-mean \
    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \
    actor_rollout_ref.rollout.enable_chunked_prefill=True \
    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \
    actor_rollout_ref.rollout.temperature=1.0 \
    actor_rollout_ref.rollout.top_p=1.0 \
    actor_rollout_ref.rollout.top_k=-1 \
    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \
    actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \
    actor_rollout_ref.rollout.val_kwargs.top_k=-1 \
    actor_rollout_ref.rollout.val_kwargs.do_sample=True \
    actor_rollout_ref.rollout.val_kwargs.n=1 \
    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \
    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \
    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \
    reward_model.reward_manager=medmatch \
    reward_model.overlong_buffer.enable=${enable_overlong_buffer} \
    reward_model.overlong_buffer.len=${overlong_buffer_len} \
    reward_model.overlong_buffer.penalty_factor=${overlong_penalty_factor} \
    trainer.logger='["console","swanlab"]' \
    trainer.project_name="${project_name}" \
    trainer.experiment_name="${exp_name}" \
    trainer.n_gpus_per_node="${gpu_num}" \
    trainer.nnodes=1 \
    trainer.val_before_train=True \
    trainer.test_freq=5 \
    trainer.save_freq=5 \
    trainer.total_epochs=5 \
    trainer.default_local_dir="${CKPTS_DIR}" \
    trainer.save_freq=500 \
    trainer.resume_mode=auto

# ===================== end =====================
